"""
Example demonstrating how to use ForwardStrategy with DIL_Frame.
"""

import logging
import torch
import torch.nn as nn
from typing import List, cast

from learners.strategies.forward_strategies import (
    create_forward_strategy,
    FORWARD_STRATEGY_REGISTRY,
    StandardForwardStrategy,
    FeatureDropoutForwardStrategy,
)


def example_usage() -> None:
    """Demonstrate how to use forward strategies."""

    # List available strategies
    available_strategies: List[str] = sorted(
        FORWARD_STRATEGY_REGISTRY.list_registered()
    )
    print(f"Available forward strategies: {available_strategies}")

    # Explicitly annotate the return types to help the type checker
    standard_strategy: StandardForwardStrategy = cast(
        StandardForwardStrategy, create_forward_strategy("standard")
    )

    # Add explicit type annotation to ensure proper typing
    dropout_strategy: FeatureDropoutForwardStrategy = (
        FORWARD_STRATEGY_REGISTRY.create_typed(
            "feature_dropout", FeatureDropoutForwardStrategy, dropout_rate=0.3
        )
    )
    # dropout_strategy: FeatureDropoutForwardStrategy = cast(
    #     FeatureDropoutForwardStrategy,
    #     create_forward_strategy("feature_dropout", dropout_rate=0.3)
    # )
    cast(FeatureDropoutForwardStrategy, dropout_strategy)

    # Now Python's type system should recognize these correctly
    print("Successfully created forward strategies:")
    print(f" - {standard_strategy.__class__.__name__}")
    # FIXME type checker doesn't recognize this as a FeatureDropoutForwardStrategy
    print(
        f" - {dropout_strategy.__class__.__name__} (dropout_rate={dropout_strategy.dropout_rate})"  # type: ignore
    )

    # Example of how this would be used with a model
    class SimpleModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(10, 2)

        def forward(self, x, training=False, ca=False):
            mode = "training" if training else ("ca" if ca else "testing")
            print(f"Model running in {mode} mode")
            return {"logits": self.fc(x)}

    # Create a model and dummy input
    model = SimpleModel()
    inputs = torch.randn(5, 10)  # Batch size 5, feature dim 10

    # Use the strategies
    print("\nTesting standard strategy:")
    _ = standard_strategy.forward_train(model, inputs)
    _ = standard_strategy.forward_test(model, inputs)
    _ = standard_strategy.forward_ca(model, inputs)

    print("\nTesting feature dropout strategy:")
    _ = dropout_strategy.forward_train(model, inputs)
    _ = dropout_strategy.forward_test(model, inputs)
    _ = dropout_strategy.forward_ca(model, inputs)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    example_usage()
